package aes

import "bytes"

type PaddingMode string

const (
	ZeroPadding  PaddingMode = "ZERO_PADDING"
	Pkcs5Padding PaddingMode = "PKCS5_PADDING"
	Pkcs7Padding PaddingMode = "PKCS7_PADDING"
)

func pkcsUnPadding(text []byte) []byte {
	n := len(text)
	if n == 0 {
		return text
	}
	paddingSize := int(text[n-1])
	return text[:n-paddingSize]
}

func pkcs7Padding(cipherText []byte, blockSize int) []byte {
	padding := blockSize - len(cipherText)%blockSize
	padText := bytes.Repeat([]byte{byte(padding)}, padding)
	return append(cipherText, padText...)
}

func pkcs5Padding(cipherText []byte) []byte {
	return pkcs7Padding(cipherText, 8)
}

func zeroPadding(in []byte, size int) []byte {
	l := 8
	switch size {
	case 16:
		l = 16
	case 24:
		l = 24
	case 32:
		l = 32
	}

	length := len(in)
	if length%l == 0 {
		return in
	} else {
		blockCount := length / l
		out := make([]byte, (blockCount+1)*l)
		var i int
		for i = 0; i < length; i++ {
			out[i] = in[i]
		}
		return out
	}
}

func zeroUnPadding(in []byte) []byte {
	return bytes.TrimRight(in, string([]byte{0}))
}
